﻿#include <QDebug>

#include "common/AVAudioConvert.h"

extern "C"
{
#include "libavcodec/avcodec.h"
#include "libavformat/avformat.h"
#include "libswresample/swresample.h"
#include "libavutil/opt.h"
};

#include "AVFileStoreAAC.h"


AVFileStoreAAC::AVFileStoreAAC(const QString &fileName) :AVFileStoreAbstract(), m_fileName(fileName)
{
    m_defaultSampleRate = 44100;
    m_defaultChannels = 2;
    m_defaultFormat = (int)AV_SAMPLE_FMT_FLTP;
}

bool AVFileStoreAAC::init()
{
    if (!m_fileName.endsWith("aac")) {
        qDebug() << "file error";
        return false;
    }
    /* 初始化容器及格式上下文 */
    avformat_alloc_output_context2(&m_formatCtx, NULL, NULL, m_fileName.toStdString().c_str());
    m_fmt = m_formatCtx->oformat;

    /* 打开待保存文件并创建关联 */
    if (avio_open(&m_formatCtx->pb,m_fileName.toStdString().c_str(), AVIO_FLAG_READ_WRITE) < 0){
        qDebug("Failed to open output file!\n");
        return false;
    }

    m_audio_st = avformat_new_stream(m_formatCtx, 0);
    if (m_audio_st == NULL){
        return false;
    }

    m_codec = avcodec_find_encoder(m_fmt->audio_codec);
    if (!m_codec){
        qDebug("Can not find encoder!\n");
        return false;
    }

    m_codecCtx = m_audio_st->codec;
    m_codecCtx->codec_id = m_fmt->audio_codec;
    m_codecCtx->codec_type = AVMEDIA_TYPE_AUDIO;
    m_codecCtx->sample_fmt = m_codec->sample_fmts[0];
    m_codecCtx->sample_rate= m_defaultSampleRate;
    m_codecCtx->channel_layout=AV_CH_LAYOUT_STEREO;
    m_codecCtx->channels = av_get_channel_layout_nb_channels(m_codecCtx->channel_layout);
    m_codecCtx->bit_rate = 128000;
    m_codecCtx->profile=FF_PROFILE_AAC_MAIN ;
    m_codecCtx->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;

    //Show some information
    av_dump_format(m_formatCtx, 0, m_fileName.toStdString().c_str(), 1);

    if (avcodec_open2(m_codecCtx, m_codec,NULL) < 0){
        qDebug("Failed to open encoder!\n");
        return false;
    }

    int size = av_samples_get_buffer_size(NULL, m_codecCtx->channels,m_codecCtx->frame_size,m_codecCtx->sample_fmt, 1);

    avformat_write_header(m_formatCtx,NULL);

    if( av_new_packet(&m_pkt,size) != 0) {
        qDebug() << "new packet error";
        return false;
    }

    return true;
}

bool AVFileStoreAAC::fillFrame(AVFrame *frame)
{
    static int i = 0;

    uint8_t dstBuf[8192] = {0}; //1024 * 4 * 2
    int retByteSize = 0;
    int retResampleSize = 0;

    AVFrame * resampleFrame = 0;

    if (frame->sample_rate != m_defaultSampleRate ||
            frame->channels != m_defaultChannels ||
            frame->format != m_defaultFormat) {
        if (resample(frame,dstBuf,&retByteSize,&retResampleSize) ) {
            resampleFrame = av_frame_alloc();
            resampleFrame->sample_rate = m_defaultSampleRate;
            resampleFrame->channels = m_defaultChannels;
            resampleFrame->format = m_defaultFormat;
            resampleFrame->nb_samples = retResampleSize;
            if (avcodec_fill_audio_frame(resampleFrame,m_defaultChannels,(AVSampleFormat)m_defaultFormat,dstBuf,retByteSize,1) < 0) {
                qDebug() << "fill error";
            }

        } else {
            qDebug() << "resample error";
        }
    } else {
        resampleFrame = av_frame_clone(frame);
    }

    resampleFrame->pts = i;
    i++;
    int gotFrame = 0;

    int ret = avcodec_encode_audio2(m_codecCtx,&m_pkt,resampleFrame,&gotFrame);

    av_frame_unref(resampleFrame);
    av_frame_free(&resampleFrame);

    if (ret < 0) {
        qDebug() << "encode error";
        return false;
    }

    if (gotFrame == 1) {
        m_pkt.stream_index = m_audio_st->index;
        ret = av_write_frame(m_formatCtx,&m_pkt);
        av_free_packet(&m_pkt);
    }

    /* 编码会用到frame空间吗？ */
    return true;
}

bool AVFileStoreAAC::destory()
{
    flush_encoder(m_formatCtx,0);

    if (m_audio_st){
        avcodec_close(m_audio_st->codec);
    }
    avio_close(m_formatCtx->pb);
    avformat_free_context(m_formatCtx);

    return true;
}

int AVFileStoreAAC::flush_encoder(AVFormatContext *fmt_ctx,unsigned int stream_index)
{
    int ret;
    int got_frame;
    AVPacket enc_pkt;
    if (!(fmt_ctx->streams[stream_index]->codec->codec->capabilities &
          AV_CODEC_CAP_DELAY))
        return 0;
    while (1) {
        enc_pkt.data = NULL;
        enc_pkt.size = 0;
        av_init_packet(&enc_pkt);
        ret = avcodec_encode_audio2 (fmt_ctx->streams[stream_index]->codec, &enc_pkt,
                                     NULL, &got_frame);
        av_frame_free(NULL);
        if (ret < 0)
            break;
        if (!got_frame){
            ret=0;
            break;
        }
        qDebug("Flush Encoder: Succeed to encode 1 frame!\tsize:%5d\n",enc_pkt.size);
        /* mux encoded frame */
        ret = av_write_frame(fmt_ctx, &enc_pkt);
        if (ret < 0)
            break;
    }

    if (ret < 0) {
        qDebug() << "flush error";
        return -1;
    }

    av_write_trailer(m_formatCtx);
    return 0;
}
